package com.pestdetection;

import org.deeplearning4j.nn.multilayer.MultiLayerNetwork;
import org.nd4j.linalg.dataset.api.iterator.DataSetIterator;
import org.nd4j.linalg.factory.Nd4j;
import java.util.ArrayList;
import java.util.Collections;
import java.util.List;

/**
 * 病虫害检测系统的主类
 * 这个程序使用深度学习来训练一个病虫害检测模型
 * 主要功能：
 * 1. 配置内存和线程
 * 2. 加载训练数据
 * 3. 创建和训练模型
 * 4. 保存训练好的模型
 */
public class Main {
    public static void main(String[] args) {
        try {
            // 第一步：配置内存和线程
            // 设置自动垃圾回收窗口为5000毫秒：可以便于无用的对象回收，节省存储空间
            Nd4j.getMemoryManager().setAutoGcWindow(5000);
            // 设置最大主线程数为4：可以提高程序的运行效率
            Nd4j.getEnvironment().setMaxMasterThreads(4);

            // 第二步：设置训练参数
            String 训练数据路径 = "dataset/train";  // 训练图片所在的文件夹
            int 图片通道数 = 3;    // RGB彩色图片
            int 训练轮数 = 10;     // 训练10轮
            int 图片高度 = 224;    // 图片高度
            int 图片宽度 = 224;    // 图片宽度

            // 第三步：加载训练数据
            System.out.println("正在加载训练数据...");
            DataSetIterator 训练数据 = DataLoader.loadData(训练数据路径);
            
            // 获取数据集中有多少种病虫害
            int 病虫害种类数 = 训练数据.getLabels().size();
            System.out.println("检测到 " + 病虫害种类数 + " 种病虫害");
            System.out.println("使用图片尺寸: " + 图片宽度 + "x" + 图片高度);

            // 打印所有病虫害种类，并按字母顺序排序，且不影响代码功能
            System.out.println("\n病虫害种类：");
            List<String> 病虫害列表 = new ArrayList<>(训练数据.getLabels());
            Collections.sort(病虫害列表);  // 按字母顺序排序
            for (String 病虫害名称 : 病虫害列表) {
                System.out.println("- " + 病虫害名称);
            }

            // 第四步：创建模型
            System.out.println("\n正在创建病虫害检测模型...");
            PestDetectionModel 模型 = new PestDetectionModel(病虫害种类数, 图片高度, 图片宽度, 图片通道数);
            
            // 第五步：训练模型
            System.out.println("开始训练模型...");
            模型.train(训练数据, 训练轮数);

            // 第六步：保存训练好的模型
            System.out.println("正在保存模型...");
            MultiLayerNetwork 训练好的模型 = 模型.getModel();
            训练好的模型.save(new java.io.File("pest_detection_model.zip"), false);

            System.out.println("模型训练完成！");
        } catch (Exception e) {
            // 如果发生错误，打印详细的错误信息
            System.err.println("发生错误: " + e.getMessage());
            e.printStackTrace();
            
            // 打印系统信息，帮助诊断问题
            System.err.println("\n当前系统配置：");
            System.err.println("- JVM最大内存: " + Runtime.getRuntime().maxMemory() / (1024*1024) + "MB");
            System.err.println("- CPU核心数: " + Runtime.getRuntime().availableProcessors());
            System.err.println("- 操作系统: " + System.getProperty("os.name"));
            System.err.println("- Java版本: " + System.getProperty("java.version"));
        }
    }
} 